import unittest

from psychsim.pwl import *
from psychsim.action import Action,ActionSet
from psychsim.world import World,stateKey
from psychsim.agent import Agent
from psychsim.reward import *

''' 
This example is used to test the dynamic turn taking
with two self.agents taking turns one after the other
'''

class TestCase(unittest.TestCase):

    def setUp(self):
        #setup world, agent, and basics
        self.maxRounds=8
        self.world = World()
        self.ping = Agent('Ping')
        self.world.addAgent(self.ping)
        self.pong = Agent('Pong')
        self.world.addAgent(self.pong)
        
        self.agents = [self.ping, self.pong]
                                    
        self.world.setOrder(['Ping', 'Pong'])

    def addStates(self):
        #define states
        self.world.defineState(None,'round',int)
        self.world.setState(None,'round',0)        
                    
    def addTermination(self):
        #terminate if maxRounds reached
        self.world.addTermination(makeTree({'if': thresholdRow(stateKey(None,'round'),self.maxRounds),
                                   True: True, False: False}))
        
    def addActions(self):
        #add actions and appropriate dynamics
        for me in self.agents: 
            action =  me.addAction({'verb': 'do nothing'})
            tree = makeTree(incrementMatrix(stateKey(None,'round'),1)) 
            self.world.setDynamics(stateKey(None,'round'),action,tree) 
            new_position_in_queue = len(self.agents)-1
            tree = makeTree(setToConstantMatrix(turnKey(me.name), new_position_in_queue )) 
            self.world.setTurnDynamics(me.name, action, tree)
        
    def testSetup(self):
        self.addStates()
        self.addTermination()
        self.addActions()

    def testRun(self):
        self.testSetup()
        self.world.printState()
        for t in range(self.maxRounds + 1):
            if t%2 == 0:
                self.assertEqual(self.world.state.domain()[0][turnKey('Ping')], 0)
                self.assertEqual(self.world.state.domain()[0][turnKey('Pong')], 1)
            else:
                self.assertEqual(self.world.state.domain()[0][turnKey('Pong')], 0)
                self.assertEqual(self.world.state.domain()[0][turnKey('Ping')], 1)
            
            self.world.explain(self.world.step(),level=1)
            self.world.state.select()
            self.world.printState()
            if self.world.terminated():
                break 

if __name__ == "__main__":
    suite = unittest.TestSuite()
    for method in dir(TestCase):
       if method.startswith("test"):
          suite.addTest(TestCase(method))
    unittest.TextTestRunner().run(suite)
            